# 测试智源研究院知识增强功能
from answer_generator import AnswerGenerator
from knowledge_enhancer import zhiyuan_enhancer

def test_knowledge_enhancement():
    print("=== 测试智源研究院知识增强功能 ===")
    
    # 测试知识统计
    print("\n1. 测试知识统计:")
    stats = zhiyuan_enhancer.get_knowledge_statistics()
    print(f"论文数量: {stats['papers_count']}")
    print(f"数据集数量: {stats['datasets_count']}")
    print(f"模型数量: {stats['models_count']}")
    print(f"新闻数量: {stats['news_count']}")
    
    # 测试知识搜索
    print("\n2. 测试知识搜索:")
    result = zhiyuan_enhancer.search_rust_knowledge("lifetime", "definition")
    print(f"搜索结果数量: {len(result.get('knowledge', []))}")
    
    # 测试回答生成增强
    print("\n3. 测试回答生成增强:")
    ag = AnswerGenerator()
    test_question = "What is lifetime?"
    test_classification = {'label': 'definition', 'confidence': 0.8}
    
    try:
        answer = ag.generate(test_question, test_classification, [], [])
        print(f"回答长度: {len(answer)}")
        print(f"包含智源研究院内容: {'智源研究院' in answer}")
        print(f"回答预览: {answer[:200]}...")
    except Exception as e:
        print(f"回答生成失败: {e}")
    
    print("\n=== 测试完成 ===")

if __name__ == "__main__":
    test_knowledge_enhancement()
